"""
股票数据分析和筛选
"""
from datetime import datetime, timedelta
import pandas as pd
import numpy as np
from sqlalchemy import select, and_
from models.stock_models import StockDailyData
from database.db_engine import SessionLocal

def get_limit_up_ma5_stocks(days=10, ma5_threshold=0.02):
    """
    获取指定天数内涨停且上一个交易日回踩5日线的股票
    
    Args:
        days (int): 查询的天数范围
        ma5_threshold (float): 判定回踩5日线的阈值，默认2%
        
    Returns:
        list: 符合条件的股票代码列表
    """
    # 获取当前日期
    end_date = datetime.now().date()
    start_date = end_date - timedelta(days=days+10)  # 多获取10天数据用于计算MA5
    
    # 创建数据库会话
    session = SessionLocal()
    try:
        # 查询股票数据
        query = select(StockDailyData).where(
            and_(
                StockDailyData.trade_date >= start_date,
                StockDailyData.trade_date <= end_date
            )
        )
        results = session.execute(query).scalars().all()
        
        # 转换为DataFrame
        df = pd.DataFrame([{
            'ts_code': r.ts_code,
            'trade_date': r.trade_date,
            'close': r.close,
            'pct_change': r.pct_change
        } for r in results])
        
        if df.empty:
            return []
        
        # 按股票代码分组处理
        result_stocks = []
        for ts_code, group in df.groupby('ts_code'):
            # 按交易日期排序
            group = group.sort_values('trade_date')
            
            # 计算5日移动平均线
            group['ma5'] = group['close'].rolling(window=5).mean()
            
            # 找出涨停的日期（涨幅大于9.8%）
            limit_up_days = group[group['pct_change'] >= 9.8]
            
            if not limit_up_days.empty:
                # 对每个涨停日进行判断
                for _, limit_up_row in limit_up_days.iterrows():
                    # 获取涨停日期的前一个交易日数据
                    prev_day = group[group['trade_date'] < limit_up_row['trade_date']].iloc[-1]
                    
                    # 判断是否回踩5日线
                    if prev_day['ma5'] is not None:
                        price_diff = abs(prev_day['close'] - prev_day['ma5']) / prev_day['ma5']
                        if price_diff <= ma5_threshold:
                            result_stocks.append({
                                'ts_code': ts_code,
                                'limit_up_date': limit_up_row['trade_date'],
                                'prev_date': prev_day['trade_date']
                            })
        
        return result_stocks
    
    finally:
        session.close()

def analyze_stock_pattern():
    """
    分析股票形态并打印结果
    """
    results = get_limit_up_ma5_stocks()
    if not results:
        print("没有找到符合条件的股票")
        return
    
    print(f"找到 {len(results)} 只符合条件的股票：")
    for stock in results:
        print(f"股票代码: {stock['ts_code']}")
        print(f"涨停日期: {stock['limit_up_date']}")
        print(f"回踩日期: {stock['prev_date']}")
        print("-" * 50)

if __name__ == "__main__":
    analyze_stock_pattern() 